import numpy as np
from .PrePocessIm import PreProcessIm
from .Prefetcher import Prefetcher
class Dataset(object):
    """
    Args:
        final_batch: bool. The last batch may not be complete, if to abandon this batch, set final_batch=False
    """
    def __init__(self,dataset_size=None,batch_size=None,final_batch=True,shuffle=True,num_prefetch_threads=1,prng=np.random,**pre_process_im_kwargs):
        self.pre_process_im = PreProcessIm(prng=prng, **pre_process_im_kwargs)
        self.prefetcher = Prefetcher(self.get_sample,dataset_size,batch_size,final_batch=final_batch,num_threads=num_prefetch_threads)
        self.shuffle = shuffle
        self.epoch_done = True
        self.prng = prng

    def get_sample(self, ptr):
        # Get one sample to put to queue
        raise NotImplementedError

    def next_batch(self):
        '''Get one sample to put to queue.'''
        raise NotImplementedError